import cv2
from ultralytics import YOLO
import torch
import time
from collections import Counter
import math

class YoloInspector:
    """
    一个YOLO巡检任务处理器，采用空间关联方法完成任务。
    - 识别视野中最主要的目标（以仪表盘为准）作为锚点。
    - 通过计算距离，关联与锚点最近的物体。
    - 在2秒的时间窗口内收集和稳定识别结果。
    - 根据要求格式化并输出最终结果。
    """
    def __init__(self, model_path, conf_threshold=0.7):
        """
        初始化巡检器。
        :param model_path: YOLO模型文件的路径 (.pt文件)。
        :param conf_threshold: 用于过滤检测结果的置信度阈值。
        """
        # 检查设备
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        print(f"正在使用设备: {self.device}")
        
        # 加载模型
        self.model = YOLO(model_path)
        
        # 状态管理变量
        self.is_inspecting = False
        self.inspection_start_time = None
        self.INSPECTION_WINDOW_S = 2.0
        self.CONF_THRESHOLD = conf_threshold

        # 用于存储2秒窗口内的检测结果
        self.collected_colors = []
        self.collected_numbers = []
        self.collected_states = []

        # 类别名称定义
        self.COLOR_CLASSES = ["red", "orange", "blue", "green", "yellow"]
        self.NUMBER_CLASSES = ["one", "two", "three", "four", "five", "six", "seven", "eight"]
        self.STATE_CLASSES = ["down", "normal", "up"]

        # 标签到中文的映射
        self.color_map = {"red": "红色", "orange": "橙色", "blue": "蓝色", "green": "绿色", "yellow": "黄色"}
        self.number_map = {"one": "1号", "two": "2号", "three": "3号", "four": "4号", "five": "5号", "six": "6号", "seven": "7号", "eight": "8号"}
        self.state_map = {"down": "偏低", "normal": "正常", "up": "偏高"}

    def _get_center(self, box):
        """计算边界框的中心点。"""
        x1, y1, x2, y2 = box
        return (x1 + x2) / 2, (y1 + y2) / 2

    def _get_dominant_object_info(self, boxes, class_names, target_classes):
        """
        从一类检测中找出最主要的一个（基于边界框面积），并返回其完整信息。
        """
        dominant_info = None
        max_area = -1

        if not boxes:
            return None

        for i in range(len(boxes)):
            class_id = int(boxes.cls[i].item())
            class_name = class_names.get(class_id)
            
            if class_name in target_classes:
                box = boxes.xyxy[i].tolist()
                area = (box[2] - box[0]) * (box[3] - box[1])
                if area > max_area:
                    max_area = area
                    dominant_info = {'class_name': class_name, 'box': box}
        
        return dominant_info

    def _find_closest_object(self, anchor_box, all_boxes, class_names, target_classes):
        """
        在所有检测到的物体中，找到一个离锚点最近的。
        """
        anchor_center = self._get_center(anchor_box)
        closest_object_name = None
        min_dist = float('inf')

        if not all_boxes:
            return None

        for i in range(len(all_boxes)):
            class_id = int(all_boxes.cls[i].item())
            class_name = class_names.get(class_id)

            if class_name in target_classes:
                box = all_boxes.xyxy[i].tolist()
                center = self._get_center(box)
                dist = math.dist(anchor_center, center)
                
                if dist < min_dist:
                    min_dist = dist
                    closest_object_name = class_name
        
        return closest_object_name

    def process_frame(self, frame):
        """
        处理单帧图像，执行检测并根据状态机管理结果。
        """
        results = self.model.predict(frame, verbose=False, conf=self.CONF_THRESHOLD)
        
        if not results or not results[0].boxes:
            if self.is_inspecting and (time.time() - self.inspection_start_time > self.INSPECTION_WINDOW_S):
                self._finalize_and_print_results()
            return

        result = results[0]
        boxes = result.boxes
        class_names = result.names

        # 1. 找出当前帧中最主要的仪表盘作为“锚点”
        dominant_state_info = self._get_dominant_object_info(boxes, class_names, self.STATE_CLASSES)
        
        if dominant_state_info is None:
            if self.is_inspecting and (time.time() - self.inspection_start_time > self.INSPECTION_WINDOW_S):
                self._finalize_and_print_results()
            return
            
        # 2. 如果检测到了锚点，我们才开始或维持一个检测窗口
        if not self.is_inspecting:
            print("检测到目标，开始2秒钟的稳定识别...")
            self.is_inspecting = True
            self.inspection_start_time = time.time()
            self.collected_colors.clear()
            self.collected_numbers.clear()
            self.collected_states.clear()

        # 3. 在检测窗口期间，持续收集信息
        if self.is_inspecting:
            anchor_box = dominant_state_info['box']
            state_name = dominant_state_info['class_name']
            
            # 4. 基于锚点的位置，找到最近的颜色和数字
            closest_color = self._find_closest_object(anchor_box, boxes, class_names, self.COLOR_CLASSES)
            closest_number = self._find_closest_object(anchor_box, boxes, class_names, self.NUMBER_CLASSES)

            # 收集结果
            self.collected_states.append(state_name)
            if closest_color:
                self.collected_colors.append(closest_color)
            if closest_number:
                self.collected_numbers.append(closest_number)

            # 检查检测窗口是否结束
            if time.time() - self.inspection_start_time > self.INSPECTION_WINDOW_S:
                self._finalize_and_print_results()

    def _finalize_and_print_results(self):
        """
        在2秒窗口结束后，汇总结果并打印。
        """
        print("2秒识别结束，正在生成最终结果...")

        try:
            # 只有成功识别出仪表盘状态，才认为是一次有效的巡检
            if not self.collected_states:
                print("在检测窗口内未能稳定识别出任何仪表盘状态。\n")
                return

            final_state_en = Counter(self.collected_states).most_common(1)[0][0]
            final_color_en = Counter(self.collected_colors).most_common(1)[0][0] if self.collected_colors else None
            final_number_en = Counter(self.collected_numbers).most_common(1)[0][0] if self.collected_numbers else None
            
            # 翻译成中文
            final_color_zh = self.color_map.get(final_color_en, "未知颜色")
            final_number_zh = self.number_map.get(final_number_en, "未知编号")
            final_state_zh = self.state_map.get(final_state_en)
            
            # 格式化输出
            output_parts = [f"{final_color_zh}区域", final_number_zh, f"仪表盘{final_state_zh}"]
            final_output = " ".join(output_parts)

            print("\n--- 最终巡检结果 ---")
            print(final_output)
            print("-" * 22 + "\n")

        except Exception as e:
            print(f"处理结果时发生错误: {e}\n")
        
        finally:
            # 重置状态机，准备下一次检测
            self.is_inspecting = False
            self.inspection_start_time = None
            self.collected_colors.clear()
            self.collected_numbers.clear()
            self.collected_states.clear()


def main():
    """
    主函数，加载视频文件并进行逐帧巡检。
    """
    # --- 配置 ---
    # !!!重要!!! 请将此路径更改为您自己训练好的YOLOv8模型文件的实际路径。
    model_path = r"D:\python.learn\python.project\yolo_dataset\250730\runs\0_epochs100\weights\best.pt"
    
    # !!!重要!!! 请将此路径更改为您要处理的MP4视频文件的实际路径。
    video_path = r"path\to\your\video.mp4"

    try:
        inspector = YoloInspector(model_path, conf_threshold=0.7)
    except Exception as e:
        print(f"错误：无法加载模型，请检查路径 '{model_path}' 是否正确。错误信息: {e}")
        return

    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"错误：无法打开视频文件 '{video_path}'。")
        return

    print(f"\n视频文件 '{video_path}' 已打开，正在进行逐帧处理... 按 Ctrl+C 可提前退出。")

    try:
        while True:
            # 读取一帧
            ret, frame = cap.read()
            # 如果ret为False，说明视频已经播放完毕
            if not ret:
                print("\n视频处理完成。")
                break
            
            # 处理每一帧图像
            inspector.process_frame(frame)
            
            # 处理视频文件时，通常不需要延时，可以全速运行
            # time.sleep(0.01)

    except KeyboardInterrupt:
        print("\n处理已由用户手动停止。")
    finally:
        # 释放视频文件句柄
        cap.release()
        cv2.destroyAllWindows()
        print("资源已释放。")

if __name__ == "__main__":
    main()

